import os
import subprocess


def main():
    project_root = os.path.abspath(
        os.path.join(os.path.dirname(__file__), '..'))
    locales_dir = os.path.join(project_root, 'locales')
    source_dir = os.path.join(project_root, 'src')

    os.makedirs(locales_dir, exist_ok=True)

    # Languages to generate/update
    languages = ['zh-CN', 'zh-TW', 'ja-JP', 'fr-FR', 'de-DE', 'ru-RU']

    # Babel configuration
    babel_config_file = os.path.join(project_root, 'babel.cfg')
    pot_file = os.path.join(locales_dir, 'messages.pot')

    # Create babel.cfg if it doesn't exist
    if not os.path.exists(babel_config_file):
        with open(babel_config_file, 'w') as f:
            f.write('[python: **.py]\n')
            f.write('[jinja2: **.html]\n')
            f.write('extensions=jinja2.ext.autoescape,jinja2.ext.with_\n')

    # 1. Extract strings using Babel
    print("Extracting translatable strings with Babel...")
    extract_command = [
        'pybabel', 'extract',
        '-F', babel_config_file,
        '-k', '_',
        '-o', pot_file,
        source_dir
    ]
    subprocess.run(extract_command, check=True, cwd=project_root)
    print(f"Created/updated template file: {pot_file}")

    # 2. Create/update .po files for each language
    for lang in languages:
        # Our standard is to use hyphens for directory names (e.g., zh-CN).
        # However, babel's `init` command requires the locale format with an underscore (e.g., zh_CN)
        # to write the correct headers in the .po file.
        # To bypass babel's path-inference logic which causes conflicts, we use the -o flag
        # to specify the output file path directly for both `init` and `update`.
        
        lang_dir = os.path.join(locales_dir, lang)
        po_dir = os.path.join(lang_dir, 'LC_MESSAGES')
        po_file = os.path.join(po_dir, 'messages.po')
        
        os.makedirs(po_dir, exist_ok=True)

        babel_lang_format = lang.replace('-', '_')

        if not os.path.exists(po_file):
            print(f"Creating .po file for {lang}...")
            command = [
                'pybabel', 'init',
                '-i', pot_file,
                '-o', po_file,  # Specify output file directly
                '-l', babel_lang_format
            ]
        else:
            print(f"Updating .po file for {lang}...")
            command = [
                'pybabel', 'update',
                '-i', pot_file,
                '-o', po_file,  # Specify output file directly
                '-l', babel_lang_format, # Explicitly specify locale for update as well
                '--ignore-pot-creation-date' # Optional: prevents minor changes on every run
            ]
        
        subprocess.run(command, check=True)

    print("\nTranslation source files are ready.")
    print("Next steps:")
    print("1. Edit the .po files in locales/<lang>/LC_MESSAGES/ to add translations.")
    print("2. Run the compile_translations.py script to generate .mo files.")


if __name__ == "__main__":
    main()
